import torch
from torch import nn
import numpy as np
def add_noise(x,noise_type, noise_level):

    if noise_type == 'gauss':
        noisy = x + torch.normal(0, noise_level/255, x.shape, device=x.device)
        noisy = torch.clamp(noisy,0,1)

    elif noise_type == 'poiss':
        noisy = torch.poisson(noise_level * x)/noise_level

    return noisy

def getPSNR(noisy_img, clean_img):
    def mse(gt: torch.Tensor, pred:torch.Tensor)-> torch.Tensor:
        loss = torch.nn.MSELoss()
        return loss(gt,pred)
    
    noisy_img, clean_img=noisy_img.to('cpu'), clean_img.to('cpu')

    noisy_img = torch.clamp(noisy_img ,0, 1)
    MSE = mse(clean_img, noisy_img).item()+1e-10
    PSNR = 10*np.log10(1/MSE)

    return PSNR